import numpy as np
from PIL import Image

def rotate(image, depth, intrinsics, pose, mask=None, k=1):
    H, W, _ = image.shape
    image_rot, depth_rot = np.rot90(image, k), np.rot90(depth, k)
    
    mask_rot = np.rot90(mask, k) if mask is not None else None

    if k % 4 == 1: # CCW
        
        R_transform = np.array([[0,-1,0,0],[1,0,0,0],[0,0,1,0],[0,0,0,1]], dtype=np.float32)
    elif k % 4 == -1 or k % 4 == 3: # CW
        R_transform = np.array([[0,1,0,0],[-1,0,0,0],[0,0,1,0],[0,0,0,1]], dtype=np.float32)
    else:
        raise NotImplementedError
    new_pose = pose @ R_transform
    
    fx, fy, cx, cy = intrinsics[0,0], intrinsics[1,1], intrinsics[0,2], intrinsics[1,2]
    fx_new, fy_new = fy, fx
    if k % 4 == 1:
        cx_new, cy_new = cy, W - cx - 1
    else:
        cx_new, cy_new = H - cy - 1, cx
    intrinsics_new = np.array([[fx_new,0,cx_new],[0,fy_new,cy_new],[0,0,1]], dtype=np.float32)
    return image_rot, depth_rot, intrinsics_new, new_pose, mask_rot

def rotate_bak(image, depth, intrinsics, pose, mask=None):
    """
    Correctly rotates an image and its associated data by 90 degrees counter-clockwise.

    This simplified function is hardcoded for k=1 (a single 90-degree CCW rotation).
    It adjusts the camera intrinsics and pose to match the new orientation.

    Args:
        image (PIL.Image or np.ndarray): The input RGB image.
        depth (np.ndarray): The input depth map.
        intrinsics (np.ndarray): A 3x3 camera intrinsics matrix.
        pose (np.ndarray): A 4x4 camera-to-world pose matrix.
        mask (np.ndarray, optional): An optional mask to rotate. Defaults to None.

    Returns:
        tuple: A tuple containing:
            - image_rot (np.ndarray): The rotated RGB image.
            - depth_rot (np.ndarray): The rotated depth map.
            - intrinsics_new (np.ndarray): The adjusted 3x3 intrinsics matrix.
            - pose_new (np.ndarray): The adjusted 4x4 pose matrix.
            - mask_rot (np.ndarray or None): The rotated mask, if provided.
    """
    if not isinstance(image, np.ndarray):
        image = np.array(image)
    
    # Get original image dimensions
    H, W, _ = image.shape

    # --- 1. Rotate Image, Depth, and Mask by 90 degrees CCW ---
    # np.rot90 with k=1 performs a single counter-clockwise rotation.
    image_rot = np.rot90(image, k=1)
    depth_rot = np.rot90(depth, k=1)
    mask_rot = np.rot90(mask, k=1) if mask is not None else None

    # --- 2. Adjust Camera Intrinsics for 90-degree CCW rotation ---
    fx, fy, cx, cy = intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2]
    
    # Swap focal lengths
    fx_new, fy_new = fy, fx
    
    # Adjust principal point
    cx_new = cy
    cy_new = W - 1 - cx
    
    intrinsics_new = np.array([
        [fx_new, 0, cx_new],
        [0, fy_new, cy_new],
        [0, 0, 1]
    ], dtype=np.float32)

    # --- 3. Adjust Camera Pose for 90-degree CCW rotation ---
    # To match a CCW rotation of the image, the camera's local coordinate system
    # must also be rotated CCW. This requires post-multiplication.
    
    # 4x4 rotation matrix for 90-degree COUNTER-clockwise rotation around the camera's local z-axis
    R_z_90_ccw = np.array([
        [ 0, -1,  0,  0],
        [ 1,  0,  0,  0],
        [ 0,  0,  1,  0],
        [ 0,  0,  0,  1]
    ], dtype=np.float32)

    # Correctly apply the rotation to the camera's local frame
    pose_new = pose @ R_z_90_ccw

    return image_rot, depth_rot, intrinsics_new, pose_new, mask_rot

